import os

import argparse
import time
import numpy as np
from matplotlib import pyplot as plt
from sklearn.metrics import auc, precision_recall_curve, roc_curve

save_dir = r'..\results\res'

def computer(human_og, human_sample_mean, human_sample_var, LLMs_og, LLMs_sample_mean, LLMs_sample_var):
    """
    计算Z-score

    @ 参数 human_og: 人类原始样本的对数似然值 List[一维数组 (片段数)] 列表长度为样本个数
    @ 参数 human_sample_mean: 人类每一片段各扰动样本的均值 List[一维数组 (片段数)]
    @ 参数 human_sample_var: 人类每一片段各扰动样本的方差 List[一维数组 (片段数)]
    @ 参数 LLMs_og: LLMs原始样本的对数似然值 List[一维数组 (片段数)]
    @ 参数 LLMs_sample_mean: LLMs每一片段各扰动样本的均值 List[一维数组 (片段数)]
    @ 参数 LLMs_sample_var: LLMs每一片段各扰动样本的方差 List[一维数组 (片段数)]

    @ 返回值 human_Z_scores: List[一维数组 (片段数) -> 该样本该片段的Z分数]
    @ 返回值 LLMs_Z_scores: List[一维数组 (片段数) -> 该样本该片段的Z分数]
    @ 返回值 human_Z_seg_var: 一维数组 (样本数) -> 该样本的Z分数方差
    @ 返回值 LLMs_Z_seg_var: 一维数组 (样本数) -> 该样本的Z分数方差

    """
    ## Step 2: 计算Z-score
    human_Z_scores = []      # 人类样本的Z-score List[一维数组 (片段数) -> 该样本该片段的Z分数]
    LLMs_Z_scores = []       # LLMs样本的Z-score List[一维数组 (片段数) -> 该样本该片段的Z分数]

    # 遍历每一个样本 计算该样本每个片段对应的Z-score
    for idx, og_ll in enumerate(human_og):
        # 若该样本中对应片段扰动的方差为0, 则将其设置为1e-6
        for idxx, val in enumerate(human_sample_var[idx]):
            if val == 0:
                human_sample_var[idx][idxx] = 1e-6
                print(f"human_sample_var[{idx}][{idxx}] is 0, set to 1e-6")

        Z_score = (og_ll - human_sample_mean[idx]) / human_sample_var[idx]
        human_Z_scores.append(Z_score)

    for idx, og_ll in enumerate(LLMs_og):
        # 若该样本中对应片段扰动的方差为0, 则将其设置为1e-6
        for idxx, val in enumerate(LLMs_sample_var[idx]):
            if val == 0:
                LLMs_sample_var[idx][idxx] = 1e-6
                print(f"LLMs_sample_var[{idx}][{idxx}] is 0, set to 1e-6")

        Z_score = (og_ll - LLMs_sample_mean[idx]) / LLMs_sample_var[idx]
        LLMs_Z_scores.append(Z_score)

    ## Step 3: 计算Z-score的方差
    human_Z_seg_var = np.array([np.var(z) for z in human_Z_scores]) # 人类样本的Z-score方差 一维数组 (样本数) -> 该样本的Z分数方差
    LLMs_Z_seg_var = np.array([np.var(z) for z in LLMs_Z_scores]) # LLMs样本的Z-score方差 一维数组 (样本数) -> 该样本的Z分数方差

    # 返回计算结果
    return human_Z_seg_var, LLMs_Z_seg_var

def load_computer(load_dir):
    """
    加载文件并调用函数计算Z-score 

    @ 参数 load_dir: str, 加载文件的根目录

    @ 返回值 human_Z_scores: List[一维数组 (片段数) -> 该样本该片段的Z分数]
    @ 返回值 LLMs_Z_scores: List[一维数组 (片段数) -> 该样本该片段的Z分数]
    @ 返回值 human_Z_seg_var: 一维数组 (样本数) -> 该样本的Z分数方差
    @ 返回值 LLMs_Z_seg_var: 一维数组 (样本数) -> 该样本的Z分数方差

    """
    human_load_dir = os.path.join(load_dir, args.dataset, "human", args.scoring_model_name)
    LLMs_load_dir = os.path.join(load_dir, args.dataset, "LLMs", args.scoring_model_name)
    ## Step 0: 加载文件
    # 需加载文件的列表
    human_og_files = [os.path.join(human_load_dir, file) for file in os.listdir(human_load_dir) if file.endswith("og.npy")]
    human_sample_files =  [os.path.join(human_load_dir, file) for file in os.listdir(human_load_dir) if file.endswith("sample.npy")]
    LLMs_og_files = [os.path.join(LLMs_load_dir, file) for file in os.listdir(LLMs_load_dir) if file.endswith("og.npy") and file.startswith(args.source_model_name)]
    LLMs_sample_files = [os.path.join(LLMs_load_dir, file) for file in os.listdir(LLMs_load_dir) if file.endswith("sample.npy") and file.startswith(args.source_model_name)]
    ## Step 1: 获取文件中的对数似然值 并计算均值和标准差
    human_og_lls = []       # 人类原始样本的对数似然值 List[一维数组 (片段数)] 列表长度为样本个数
    LLMs_og_lls = []        # LLMs原始样本的对数似然值 List[一维数组 (片段数)]

    human_seg_means = []     # 人类每一片段各扰动样本的均值 List[一维数组 (片段数)]
    huamn_seg_var = []      # 人类每一片段各扰动样本的方差 List[一维数组 (片段数)]
    LLMs_seg_means = []      # LLMs每一片段各扰动样本的均值 List[一维数组 (片段数)]
    LLMs_seg_var = []       # LLMs每一片段各扰动样本的方差 List[一维数组 (片段数)]
    # 获取human的数据
    for file in human_og_files:
        human_og_lls.append(np.load(file).squeeze())
    for file in human_sample_files:
        sample = np.load(file)
        # sample: 二维数组 (扰动样本数, 片段数)
        human_seg_means.append(np.mean(sample, axis=0))
        huamn_seg_var.append(np.std(sample, axis=0))
    # 获取LLMs的数据
    for file in LLMs_og_files:
        LLMs_og_lls.append(np.load(file).squeeze())
    for file in LLMs_sample_files:
        sample = np.load(file)
        LLMs_seg_means.append(np.mean(sample, axis=0))
        LLMs_seg_var.append(np.std(sample, axis=0))
    # 返回计算结果
    return computer(human_og_lls, human_seg_means, huamn_seg_var, LLMs_og_lls, LLMs_seg_means, LLMs_seg_var)

def plot_hist(human, LLM):
    """
    输出结果

    绘制直方图
    """
    # 计算Z-score的均值和方差
    human_Z_var_mean = np.mean(human_Z_seg_var, axis=0)  # 人类样本的片段Z-score方差的均值
    human_Z_var_var = np.var(human_Z_seg_var, axis=0)  # 人类样本的片段Z-score方差的方差
    LLMs_Z_var_mean = np.mean(LLMs_Z_seg_var, axis=0)  # LLMs样本的片段Z-score方差的均值
    LLMs_Z_var_var = np.var(LLMs_Z_seg_var, axis=0)  # LLMs样本的片段Z-score方差的方差
    print("|+------------------------------+|")
    print("||           RESULTS            ||")
    print(f"| human mean of variances  : {human_Z_var_mean:.3f} |")
    print(f"| human var of variances   : {human_Z_var_var:.3f} |")
    print(f"| LLM mean of variances: {LLMs_Z_var_mean:.3f} |")
    print(f"| LLM var of variances : {LLMs_Z_var_var:.3f} |")
    print("+--------------------------------+")

    # 计算roc_auc
    fpr, tpr, roc_thresholds = roc_curve([1] * len(human_Z_seg_var) + [0] * len(LLMs_Z_seg_var), np.concatenate((human_Z_seg_var, LLMs_Z_seg_var)))
    roc_auc = auc(fpr, tpr)
    # 计算pr_auc
    precision, recall, pr_thresholds = precision_recall_curve([1] * len(human_Z_seg_var) + [0] * len(LLMs_Z_seg_var), np.concatenate((human_Z_seg_var, LLMs_Z_seg_var)))
    pr_auc = auc(recall, precision)
    #计算分类阈值
    youden_j = tpr - fpr
    best_threshold_index = np.argmax(youden_j)
    best_threshold = roc_thresholds[best_threshold_index]

    # 输出结果
    print(f"ROC_AUC: {roc_auc:.3f}, PR_AUC: {pr_auc:.3f}")
    print(f"Best threshold based on ROC: {best_threshold:.3f}")

    # 将结果保存到文件中
    save_path = os.path.join(save_dir, args.dataset, f"{start_time}_{args.source_model_name}_{args.scoring_model_name}", "auc.txt")
    if not os.path.exists(os.path.dirname(save_path)):
        os.makedirs(os.path.dirname(save_path))
    with open(save_path, 'w') as f:
        f.write(f"ROC_AUC: {roc_auc:.3f}, PR_AUC: {pr_auc:.3f}, Best threshold based on ROC: {best_threshold:.3f}")
    
    plt.figure()
    bins = np.linspace(min(np.min(human), np.min(LLM)), max(np.max(human), np.max(LLM)), 70+1)
    plt.hist(human, bins=bins, color='blue', alpha=0.7, label='human')
    plt.hist(LLM, bins=bins, color='orange', alpha=0.7, label='LLM')
    plt.legend()

    # 将图像保存到对应的文件夹
    save_path = os.path.join(save_dir, args.dataset, f"{start_time}_{args.source_model_name}_{args.scoring_model_name}", "hist.png")
    if not os.path.exists(os.path.dirname(save_path)):
        os.makedirs(os.path.dirname(save_path))
    plt.savefig(save_path)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--dataset', type=str, default="xsum")  # 使用数据集名称
    parser.add_argument('--source_model_name', type=str, required=True, help='Please specify the source model name.')  # 生成模型名称
    parser.add_argument('--scoring_model_name', type=str, required=True, help='Please specify the scoring model name.')  # 计算得分模型名称
    args = parser.parse_args()

    # 获取当前时间
    start_time =  time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime())
    
    load_dir = r"..\results\score"

    # 加载数据并计算    
    human_Z_seg_var, LLMs_Z_seg_var = load_computer(load_dir)

    # 计算AUROC 并将全部样本的各Z-score方差绘制直方图
    plot_hist(human_Z_seg_var, LLMs_Z_seg_var)